#include "tls_socket_operate.h"

#include "./../description_module/description_module.h"

#include <cstdint>
#include <cstdio>
#include <fcntl.h>
#include <memory>
#include <openssl/bio.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <pthread.h>
#include <string>
#include <unistd.h>
#include <utility>
#include <cstring>

SSL_CTX *         TlsSocketOperate::s_m_client_ssl_ctx = nullptr;
SSL_CTX *         TlsSocketOperate::s_m_server_ssl_ctx = nullptr;
SSL_CTX *         TlsSocketOperate::s_m_udp_client_ssl_ctx = nullptr;
SSL_CTX *         TlsSocketOperate::s_m_udp_server_ssl_ctx = nullptr;
std::string       TlsSocketOperate::s_m_ssl_log_file;
pthread_once_t    TlsSocketOperate::s_m_once           = PTHREAD_ONCE_INIT;

TlsSocketOperate::TlsSocketOperate(std::shared_ptr<Description> description, int operate_type)
{
    pthread_once(&s_m_once, []() { InitTlsSocket(); });
    m_bio_ssl      = nullptr;
    m_description  = std::move(description);
    m_operate_type = operate_type;
}

TlsSocketOperate::~TlsSocketOperate()
{
    Close();
}

void TlsSocketOperate::InitTlsSocket()
{
    SSL_library_init();
    s_m_client_ssl_ctx = SSL_CTX_new(TLS_client_method());
    s_m_server_ssl_ctx = SSL_CTX_new(TLS_server_method());
    s_m_udp_client_ssl_ctx = SSL_CTX_new(DTLS_client_method());
    s_m_udp_server_ssl_ctx = SSL_CTX_new(DTLS_server_method());
}

int TlsSocketOperate::Connection()
{
    return Connection(TlsSocketOperate::TlsTcpOption{});
}

int TlsSocketOperate::Connection(const TlsSocketOperate::TlsTcpOption& option)
{
    m_bio_ssl = CreateSSL(m_operate_type, m_description);
    if (m_bio_ssl)
    {
        int result = ConfigSSL(m_bio_ssl, option);
        if (result >= 0)
        {
            result = BIO_do_connect(m_bio_ssl);
            if (result < 0)
            {
                result = -1;
                if (BIO_should_retry(m_bio_ssl))
                {
                    result = 0;
                }
            }
        }
        return result;
    }
    return -1;
}

int TlsSocketOperate::ServerConnection()
{
    return ServerConnection(TlsSocketOperate::TlsTcpOption{});
}

int TlsSocketOperate::ServerConnection(const TlsTcpOption & option)
{
    if (m_bio_ssl)
    {
        return HandShake();
    }
    else
    {
        m_bio_ssl = CreateSSL(m_operate_type, m_description);
        if (m_bio_ssl)
        {
            int result = ConfigSSL(m_bio_ssl, option);
            if (result >= 0)
            {
                result = BIO_do_accept(m_bio_ssl);
                if (result < 0)
                {
                    result = -1;
                    if (BIO_should_retry(m_bio_ssl))
                    {
                        result = 0;
                    }
                }
            }
            return result;
        }
        return -1;
    }
}
int TlsSocketOperate::HandShake() const
{
    int result = BIO_do_handshake(m_bio_ssl);
    if (result < 0)
    {
        if (BIO_should_retry(m_bio_ssl))
        {
            result = 0;
        }
        else
        {
            result = -1;
        }
    }
    return result;
}

void TlsSocketOperate::Close()
{
    if (m_bio_ssl)
    {
        BIO_free_all(m_bio_ssl);
        m_bio_ssl = nullptr;
    }
}

int TlsSocketOperate::Write(const uint8_t * data, int length) const
{
    if (m_bio_ssl)
    {
        int result = BIO_write(m_bio_ssl, data, length);
        if (result <= 0)
        {
            result = -1;
            if (BIO_should_retry(m_bio_ssl))
            {
                result = 0;
            }
        }
        return result;
    }
    return -1;
}

int TlsSocketOperate::Read(uint8_t * data, int length) const
{
    if (m_bio_ssl)
    {
        int result = BIO_read(m_bio_ssl, data, length);
        if (result <= 0)
        {
            result = -1;
            if (BIO_should_retry(m_bio_ssl))
            {
                result = 0;
            }
        }
        return result;
    }
    return -1;
}

BIO * TlsSocketOperate::CreateSSL(int operate_type, const std::shared_ptr<Description>& description)
{
    BIO * bio = nullptr;
    switch (operate_type)
    {
        case EmOperateTls:
        {
            bio = BIO_new_ssl_connect(s_m_client_ssl_ctx);
            SSL * ssl = nullptr;
            BIO_get_ssl(bio, &ssl);
            SSL_set_fd(ssl, description->DescriptionData());
            break;
        }
        case EmOperateDTls:
        {
            bio = BIO_new_dgram(description->DescriptionData(), BIO_NOCLOSE);
            SSL *ssl = SSL_new(s_m_udp_client_ssl_ctx);
            BIO *ssl_bio = BIO_new_ssl(s_m_udp_client_ssl_ctx, 1);
            BIO_push(ssl_bio, bio);
            BIO_set_ssl(ssl_bio, ssl, BIO_NOCLOSE);
            bio = ssl_bio;
            break;
        }
        case EmOperateTlsServer:
        {
            bio = BIO_new_ssl(s_m_server_ssl_ctx, 0);
            SSL * ssl = nullptr;
            BIO_get_ssl(bio, &ssl);
            SSL_set_fd(ssl, description->DescriptionData());
            break;
        }
        case EmOperateDTlsServer:
        {
            bio = BIO_new_dgram(description->DescriptionData(), BIO_NOCLOSE);
            SSL *ssl = SSL_new(s_m_udp_server_ssl_ctx);
            BIO *ssl_bio = BIO_new_ssl(s_m_udp_server_ssl_ctx, 0);
            BIO_push(ssl_bio, bio);
            BIO_set_ssl(ssl_bio, ssl, BIO_NOCLOSE);
            bio = ssl_bio;
            break;
        }
        default:
            break;
    }
    return bio;
}

int TlsSocketOperate::ConfigSSL(BIO * bio, const TlsSocketOperate::TlsOption & option)
{
    SSL * ssl = nullptr;
    BIO_get_ssl(bio, &ssl);
    int result = 0;
    if (!option.m_private_key_file.empty())
    {
        result = SSL_use_PrivateKey_file(ssl, option.m_private_key_file.c_str(), SSL_FILETYPE_PEM);
        if (result <= 0)
        {
            return -1;
        }
    }
    if (!option.m_public_crt_file.empty())
    {
        result = SSL_use_certificate_file(ssl, option.m_public_crt_file.c_str(), SSL_FILETYPE_PEM);
        if (result <= 0)
        {
            return -1;
        }
    }
    if (!option.m_ca_file.empty())
    {
        if (!SSL_is_dtls(ssl))
        {
            if (SSL_is_server(ssl))
            {
                SSL_CTX_load_verify_locations(s_m_server_ssl_ctx, option.m_ca_file.c_str(), nullptr);
            }
            else
            {
                SSL_CTX_load_verify_locations(s_m_client_ssl_ctx, option.m_ca_file.c_str(), nullptr);
            }
        }
        else
        {
            if (SSL_is_server(ssl))
            {
                SSL_CTX_load_verify_locations(s_m_udp_server_ssl_ctx, option.m_ca_file.c_str(), nullptr);
            }
            else
            {
                SSL_CTX_load_verify_locations(s_m_udp_client_ssl_ctx, option.m_ca_file.c_str(), nullptr);
            }
        }
    }
    if (!option.m_host_name.empty())
    {
        result = SSL_set_tlsext_host_name(ssl, option.m_host_name.c_str());
        if (result <= 0)
        {
            return -1;
        }
    }
    if (ssl != nullptr && !option.m_ssl_key_log_file.empty())
    {
        SSL_CTX * ssl_ctx = SSL_get_SSL_CTX(ssl);
        if (ssl_ctx && SSL_CTX_get_keylog_callback(ssl_ctx) == nullptr)
        {
            if (s_m_ssl_log_file.empty())
            {
                s_m_ssl_log_file = option.m_ssl_key_log_file;
            }
            SSL_CTX_set_keylog_callback(ssl_ctx, WriteSSLLogFile);
        }
    }
    return result;
}

void TlsSocketOperate::WriteSSLLogFile(const SSL * ssl, const char * line)
{
    int fp = open(s_m_ssl_log_file.c_str(), O_APPEND | O_CREAT | O_RDWR, 0666);
    if (fp >= 0)
    {
        write(fp, line, strlen(line));
        write(fp, "\n", 1);
        close(fp);
    }
}

